#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import random
import argparse
import os
import numpy as np
from collections import OrderedDict
import struct
import cv2
import math
from math import cos, sin
import re
from matplotlib import pyplot as plt

# SLEEP_TIME = 40
SLEEP_TIME = 1

image_base_path = r'/rootfs/media/kasim/Data1/data/VideoCropFace'
out_path = r'/rootfs/media/kasim/Data1/data/VideoCropFace'
file_list = r'/rootfs/media/kasim/Data1/data/VideoCropFace/video_crop_face_list.attr.txt'
bbox_file_list = r'/rootfs/media/kasim/Data1/data/VideoCropFace/video_crop_face_list.txt'

# image_base_path = r'/rootfs/media/yery/Dataset/VideoCropFace'
# out_path = r'/rootfs/media/yery/Dataset/VideoCropFace'
# file_list = r'/rootfs/media/yery/Dataset/VideoCropFace/video_crop_face_list.attr.txt'
# bbox_file_list = r'/rootfs/media/yery/Dataset/VideoCropFace/video_crop_face_list.txt'


SHOW = True
DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720

h_image_count = 8
v_image_count = 4
img_width = 224
img_height = 224

min_pitch = -20
max_pitch = 20
min_roll = -30
max_roll = 30
min_yaw = -20
max_yaw = 20

min_angle_score = 0.1
max_angle_diff = 20

min_width = 64
min_height = 82


def draw_hist(plt, name, xlist, bins):
    if isinstance(xlist, np.ndarray):
        x = xlist
    else:
        x = np.array(xlist)
    n, bins, patches = plt.hist(x, bins, align=u'mid', facecolor='blue', alpha=0.5)
    plt.title('{}: {:.2f} {:.3f} {:.2f} {:.2f}'.format(name, x.mean(), x.std(), x.min(), x.max()))


def show_hist(value_list, min_value=0.0, max_value=1.0, dbin=0.1, dxticks=0.1, title_str='Age'):
    bins = np.linspace(min_value, max_value, int((max_value - min_value) / dbin) + 1)
    xticks = np.linspace(min_value, max_value, int((max_value - min_value) / dxticks) + 1)

    fig = plt.figure(title_str, figsize=(24, 12))
    plt.subplot(111)
    draw_hist(plt, title_str, value_list, bins)
    plt.xticks(xticks)
    plt.tight_layout()
    plt.show()


def show_hist_ex(value_list, bin_count=20, scale=1.0, title_str='Age', mode=0):
    if isinstance(value_list, np.ndarray):
        values = value_list
    else:
        values = np.array(value_list)
    if scale != 1.0:
        values = values * scale
    min_value = values.min()
    max_value = values.max()
    min_value = math.floor(min_value)
    max_value = math.ceil(max_value)

    diff_value = max_value - min_value
    if mode != 0:
        if bin_count > diff_value:
            bin_count = diff_value
    dxticks = diff_value / bin_count

    show_hist(values, min_value, max_value, dxticks, dxticks, title_str)


def cal_angle(bin_angle, k=5):
    # bin_angle = np.array(bin_angle, dtype=np.float32)
    bin_angle_index0 = np.argmax(bin_angle, 1)
    idxes = np.arange(0, bin_angle.shape[0])
    bin_angle_prob0 = bin_angle[idxes, bin_angle_index0]

    end_idx = bin_angle.shape[1] - 1

    bin_angle_probs = []
    bin_angle_indexes = []
    c = 1 + k // 2
    for offset in range(1, c):
        bin_angle_index1 = bin_angle_index0 - offset
        bin_angle_index2 = bin_angle_index0 + offset

        start_mask = bin_angle_index1 < 0
        end_mask = bin_angle_index2 > end_idx

        bin_angle_index1[start_mask] = 0
        bin_angle_index2[end_mask] = 0

        bin_angle_prob1 = bin_angle[idxes, bin_angle_index1]
        bin_angle_prob2 = bin_angle[idxes, bin_angle_index2]

        bin_angle_prob1[start_mask] = 0
        bin_angle_prob2[end_mask] = 0

        bin_angle_indexes.append(bin_angle_index1)
        bin_angle_indexes.append(bin_angle_index2)
        bin_angle_probs.append(bin_angle_prob1)
        bin_angle_probs.append(bin_angle_prob2)

    bin_angle_indexes.append(bin_angle_index0)
    bin_angle_probs.append(bin_angle_prob0)

    bin_angle_prob = np.stack(bin_angle_probs, 1)
    bin_angle_prob_norm = np.sum(bin_angle_prob, 1, keepdims=True)
    bin_angle_prob = bin_angle_prob / bin_angle_prob_norm

    bin_angle_index = np.stack(bin_angle_indexes, 1)
    bin_angle = 3 * bin_angle_index - 99

    bin_angle = bin_angle_prob * bin_angle
    # bin_angle = np.sum(bin_angle, 1, keepdims=True)
    bin_angle = np.sum(bin_angle, 1)

    angle_score = bin_angle_prob0

    return bin_angle, angle_score


def merge_images(image_list, h_image_count=1, v_image_count=1):
    display_image_count = h_image_count * v_image_count
    total_image_count = len(image_list)
    last_idx = total_image_count - 1

    m_images = []
    imgs = []
    for idx, img in enumerate(image_list):
        imgs.append(img)
        file_count = len(imgs)
        if file_count < display_image_count and (idx < last_idx):
            continue

        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            img_width = img.shape[1]
            img_height = img.shape[0]
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        imgs.clear()

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        m_image = np.vstack(images)
        m_images.append(m_image)

    return m_images


def read_image(image_dir, file_name, img_width, img_height, mode=1):
    file_path = os.path.join(image_dir, file_name.strip())
    img = cv2.imread(file_path)
    if mode == 0:
        img = cv2.resize(img, (img_width, img_height))
    else:
        w, h = img.shape[1], img.shape[0]
        wscale = img_width / w
        hscale = img_height / h
        scale = min(wscale, hscale)
        w = int(round(w * scale))
        h = int(round(h * scale))
        image = cv2.resize(img, (w, h))
        img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
        x = (img_width-w)//2
        y = (img_height-h)//2
        img[y:y+h, x:x+w, :] = image
    return img


def draw_axis(image, yaw, pitch, roll, tdx, tdy, size=100):
    tdx = int(round(tdx))
    tdy = int(round(tdy))

    LINE_THICKNESS = 4
    # Convert from angular to radian degree
    pitch = pitch * np.pi / 180
    yaw = -(yaw * np.pi / 180)
    roll = roll * np.pi / 180

    # X-Axis pointing to right. drawn in red, pitch绕此轴旋转
    x1 = size * (cos(yaw) * cos(roll)) + tdx
    y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy

    # Y-Axis | drawn in green, yaw绕此轴旋转
    x2 = size * (-cos(yaw) * sin(roll)) + tdx
    y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy

    # Z-Axis (out of the screen) drawn in blue, roll绕此轴旋转
    x3 = size * (sin(yaw)) + tdx
    y3 = size * (-cos(yaw) * sin(pitch)) + tdy

    x1 = int(round(x1))
    y1 = int(round(y1))
    x2 = int(round(x2))
    y2 = int(round(y2))
    x3 = int(round(x3))
    y3 = int(round(y3))

    cv2.line(image, (tdx, tdy), (x1, y1), (0, 0, 255), thickness=LINE_THICKNESS)
    cv2.line(image, (tdx, tdy), (x2, y2), (0, 255, 0), thickness=LINE_THICKNESS)
    cv2.line(image, (tdx, tdy), (x3, y3), (255, 0, 0), thickness=LINE_THICKNESS)


def show_error_images(file_infos):
    img_list = []
    for i in range(len(file_infos)):
        image_file = file_infos[i][0]
        img = read_image(image_base_path, image_file, img_width, img_height)
        img_list.append(img)

    use_bin = True
    while True:
        _img_list = []
        for img, info in zip(img_list, file_infos):
            image_file, pitch, roll, yaw, _pitch, _roll, _yaw, pitch_diff, roll_diff, yaw_diff, _pitch_score, _roll_score, _yaw_score = info
            image = img.copy()

            info_str = 'P {:.01f} {:.01f}'.format(pitch, _pitch)
            cv2.putText(image, info_str, (2, 16), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 2)

            info_str = 'Y {:.01f} {:.01f}'.format(yaw, _yaw)
            cv2.putText(image, info_str, (2, 32), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 2)

            info_str = 'R {:.01f} {:.01f}'.format(roll, _roll)
            cv2.putText(image, info_str, (2, 48), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 0, 0), 2)

            info_str = 'PS {:.03f}'.format(_pitch_score)
            cv2.putText(image, info_str, (2, 64), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 2)

            info_str = 'YS {:.03f}'.format(_yaw_score)
            cv2.putText(image, info_str, (2, 80), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 2)

            info_str = 'RS {:.03f}'.format(_roll_score)
            cv2.putText(image, info_str, (2, 96), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 0, 0), 2)

            info_str = 'PD {:.03f}'.format(pitch_diff)
            cv2.putText(image, info_str, (2, 112), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 2)

            info_str = 'YD {:.03f}'.format(yaw_diff)
            cv2.putText(image, info_str, (2, 128), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 2)

            info_str = 'RD {:.03f}'.format(roll_diff)
            cv2.putText(image, info_str, (2, 144), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 0, 0), 2)

            if use_bin:
                pitch, roll, yaw = _pitch, _roll, _yaw

            draw_axis(image, yaw, pitch, roll, tdx=img_width / 2, tdy=img_height / 2, size=img_width / 2)

            _img_list.append(image)

        images = merge_images(_img_list, h_image_count, v_image_count)
        image = images[0]
        cv2.imshow('headpose', image)
        key = cv2.waitKey()
        if key == 27 or key == ord('q') or key == ord('Q'):
            return True
        elif key == ord('e') or key == ord('E'):
            use_bin = False
        elif key == ord('b') or key == ord('B'):
            use_bin = True
        elif key == ord('n') or key == ord('N') or key == ord(' '):
            break
    return False


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    image_list = []
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_path = line.strip()
            image_list.append(file_path)

    _base_path = os.path.splitext(file_list)[0]
    title_file_path = _base_path + '.title.txt'
    with open(title_file_path, 'r') as file:
        title_str = file.readline()
        attr_item_count_str = file.readline()
        titles = title_str.strip().split(';')
        attr_item_counts = attr_item_count_str.strip().split(';')
        attr_item_counts = list(map(int, attr_item_counts))
        total_attr_item_count = sum(attr_item_counts)
        # attr_count = len(titles)

    attr_item_offsets = [0]
    offset = 0
    for attr_item_count in attr_item_counts:
        offset += attr_item_count
        attr_item_offsets.append(offset)
    attrs_size = total_attr_item_count * 4
    image_attributes_dict = {}
    data_file_path = _base_path + '.dat'

    with open(data_file_path, 'rb') as file:
        file_count = 0
        for image_name in image_list:
            dat = file.read(attrs_size)
            datas = struct.unpack('{}f'.format(total_attr_item_count), dat)

            attribute = {}
            for i, title in enumerate(titles):
                start_offset = attr_item_offsets[i]
                end_offset = attr_item_offsets[i+1]
                attribute[title] = datas[start_offset:end_offset]

            image_attributes_dict[image_name] = attribute

            file_count += 1
            # if file_count >= 1000:
            #     break

    image_bbox_dict = {}
    with open(bbox_file_list, 'r') as file:
        for line in file.readlines():
            lines = line.split()
            image_name = lines[0].strip()
            bbox = lines[1].strip()
            image_bbox_dict[image_name] = bbox

    total_file_count = len(image_attributes_dict)
    file_count = 0

    pitch_list = []
    roll_list = []
    yaw_list = []
    bin_pitch_list = []
    bin_roll_list = []
    bin_yaw_list = []
    quality_list = []
    image_names = list(image_attributes_dict.keys())
    image_names.sort()
    for image_name in image_names:
        attribute = image_attributes_dict[image_name]
        file_count += 1
        if file_count % 1000 == 0:
            print('Process File Count: {}/{}'.format(file_count, total_file_count))

        quality = attribute['quality']
        quality = np.array(quality, dtype=np.float32)
        quality_list.append(quality)

        pitch = attribute['pitch'][0]
        yaw = attribute['yaw'][0]
        roll = attribute['roll'][0]

        bin_pitch = attribute['bin_pitch']
        bin_yaw = attribute['bin_yaw']
        bin_roll = attribute['bin_roll']

        bin_pitch = np.array(bin_pitch, dtype=np.float32)
        bin_yaw = np.array(bin_yaw, dtype=np.float32)
        bin_roll = np.array(bin_roll, dtype=np.float32)

        pitch_list.append(pitch)
        roll_list.append(roll)
        yaw_list.append(yaw)

        bin_pitch_list.append(bin_pitch)
        bin_roll_list.append(bin_roll)
        bin_yaw_list.append(bin_yaw)

    print('Process File Count: {}/{}'.format(file_count, total_file_count))

    pitchs = np.array(pitch_list, dtype=np.float32)
    rolls = np.array(roll_list, dtype=np.float32)
    yaws = np.array(yaw_list, dtype=np.float32)

    bin_pitchs = np.array(bin_pitch_list, dtype=np.float32)
    bin_rolls = np.array(bin_roll_list, dtype=np.float32)
    bin_yaws = np.array(bin_yaw_list, dtype=np.float32)

    _pitchs, _pitch_scores = cal_angle(bin_pitchs)
    _rolls, _roll_scores = cal_angle(bin_rolls)
    _yaws, _yaw_scores = cal_angle(bin_yaws)

    pitch_mask = (min_pitch <= pitchs) & (pitchs <= max_pitch)
    roll_mask = (min_roll <= rolls) & (rolls <= max_roll)
    yaw_mask = (min_yaw <= yaws) & (yaws <= max_yaw)

    pitch_diffs = np.abs(pitchs - _pitchs)
    roll_diffs = np.abs(rolls - _rolls)
    yaw_diffs = np.abs(yaws - _yaws)

    pitch_diffs_mask = pitch_diffs < max_angle_diff
    roll_diffs_mask = roll_diffs < max_angle_diff
    yaw_diffs_mask = yaw_diffs < max_angle_diff

    pitch_scores_mask = _pitch_scores > min_angle_score
    roll_scores_mask = _roll_scores > min_angle_score
    yaw_scores_mask = _yaw_scores > min_angle_score

    mask = pitch_mask & roll_mask & yaw_mask & pitch_diffs_mask & roll_diffs_mask & yaw_diffs_mask & pitch_scores_mask & roll_scores_mask & yaw_scores_mask

    qualitys = np.array(quality_list, dtype=np.float32)
    # qualitys = qualitys[mask]
    qualitys = 0.3333333333 * qualitys[:, 2] + 0.66666666 * qualitys[:, 1] + qualitys[:, 0]

    out_file_path = os.path.join(out_path, 'face_image_filter.txt')
    with open(out_file_path, 'w') as file:
        for i, m in enumerate(mask):
            if m:
                image_name = image_names[i]
                quality = qualitys[i]
                bbox = image_bbox_dict.get(image_name, None)
                if bbox is None:
                    continue
                bb = [int(v) for v in bbox.split(',')]
                if ((bb[2] - bb[0]) < min_width) or ((bb[3] - bb[1]) < min_height):
                    continue
                out_info = '{} {} {}\n'.format(image_name, bbox, quality)
                file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    print('Finish!')


if __name__ == '__main__':
    main()
